import os
import cv2
import numpy as np
import torch
from random import randint
from models import yolo
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.torch_utils import select_device
from PIL import Image, ImageDraw, ImageFont

class detector(object):
    def __init__(self):
        self.img_size = 640
        self.threshold = 0.3
        self.max_frame = 160
        self.init_model()
        self.textSize = 12
        self.fontStyle = ImageFont.truetype("simsun.ttc", self.textSize, encoding="utf-8")

    def init_model(self):
        self.weights = 'weights/best.pt'
        self.device = '0' if torch.cuda.is_available() else 'cpu'
        self.device = select_device(self.device)
        m = attempt_load(self.weights, map_location=self.device)
        m.to(self.device).eval()
        m.float()
        self.model = m
        #self.names = m.module.names if hasattr(m, 'module') else m.names
        self.names = [
            '可回收物-玻璃瓶','其他垃圾-餐盒','可回收物-其他纸箱',
            '可回收物-透明塑料瓶','可回收物-塑料瓶盖','可回收物-饮料罐',
            '可回收物-食物罐','可回收物-其他塑料瓶','可回收物-拉环',
            '有害垃圾-气雾剂','可回收物-玻璃杯','可回收物-其他塑料包装',
            '可回收物-泡沫塑料片','可回收物-塑料薄膜','可回收物-其他塑料',
            '可回收物-饮料盒','可回收物-金属瓶盖','可回收物-一次性食品盒',
            '可回收物-普通纸张','可回收物-纸杯','有害垃圾-香烟',
            '可回收物-一次性手提袋','可回收物-纸巾','可回收物-马桶管',
            '可回收物-食品袋', '可回收物-塑料盖','可回收物-金属盖',
            '可回收物-鸡蛋盒','可回收物-塑料吸管','可回收物-纸袋',
            '可回收物-一次性塑料杯','可回收物-碎玻璃','可回收物-塑料器皿',
            '可回收物-玻璃罐','厨余垃圾-食物垃圾','可回收物-挤压管',
            '可回收物-饭盒','可回收物-鞋','可回收物-垃圾袋',
            '可回收物-铝箔','可回收物-金属环','可回收物-泡沫杯',
            '可回收物-纸吸管','可回收物-波纹纸箱','其他垃圾-未知杂物',
            '可回收物-铝泡罩包装','有害垃圾-电池','其他垃圾-绳索',
            '可回收物-其他塑料容器','可回收物-聚丙烯袋','可回收物-废金属',
            '可回收物-杂志纸','可回收物-比萨饼盒','可回收物-塑料手套',
            '可回收物-包装纸','可回收物-吸塑包装','可回收物-泡沫食品容器',
            '可回收物-特百惠','可回收物-其他塑料杯']
        self.colors = [(randint(0, 255), randint(0, 255), randint(0, 255)) for _ in self.names]

    def letterbox(self, img, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True):
        shape = img.shape[:2]  # current shape [height, width]
        if isinstance(new_shape, int):
            new_shape = (new_shape, new_shape)

        # Scale ratio (new / old)
        r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
        if not scaleup:  # only scale down, do not scale up (for better test mAP)
            r = min(r, 1.0)

        # Compute padding
        ratio = r, r  # width, height ratios
        new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
        dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1]  # wh padding
        if auto:  # minimum rectangle
            dw, dh = np.mod(dw, 32), np.mod(dh, 32)  # wh padding
        elif scaleFill:  # stretch
            dw, dh = 0.0, 0.0
            new_unpad = (new_shape[1], new_shape[0])
            ratio = new_shape[1] / shape[1], new_shape[0] / shape[0]  # width, height ratios

        dw /= 2  # divide padding into 2 sides
        dh /= 2

        if shape[::-1] != new_unpad:  # resize
            img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)
        top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
        left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
        img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)  # add border
        return img, ratio, (dw, dh)


    def preprocess(self, img):
        img0 = img.copy()
        img = self.letterbox(img, new_shape=640)[0]
        img = img[:, :, ::-1].transpose(2, 0, 1)
        img = np.ascontiguousarray(img)
        img = torch.from_numpy(img).to(self.device)
        img = img.float()  # 半精度
        img /= 255.0  # 图像归一化
        if img.ndimension() == 3:
            img = img.unsqueeze(0)
        return img0, img

    def cv2AddChineseText(self, img, text, position, textColor=(0, 255, 0), textSize=12):
        if (isinstance(img, np.ndarray)):  # 判断是否OpenCV图片类型
            img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        # 创建一个可以在给定图像上绘图的对象
        draw = ImageDraw.Draw(img)
        # 字体的格式
        #fontStyle = ImageFont.truetype("simsun.ttc", textSize, encoding="utf-8")
        # 绘制文本
        draw.text(position, text, textColor, font=self.fontStyle)
        # 转换回OpenCV格式
        return cv2.cvtColor(np.asarray(img), cv2.COLOR_RGB2BGR)

    def plot_bboxes(self, image, bboxes, line_thickness=None):
        tl = line_thickness or round(0.002 * (image.shape[0] + image.shape[1]) / 2) + 1  # line/font thickness
        for (x1, y1, x2, y2, cls_id, conf) in bboxes:
            color = self.colors[self.names.index(cls_id)]
            c1, c2 = (x1, y1), (x2, y2)
            cv2.rectangle(image, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)
            #tf = max(tl - 1, 1)  # font thickness
            #t_size = cv2.getTextSize(cls_id, 0, fontScale=tl / 3, thickness=tf)[0]
            #c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
            c2 = c2[0], c1[1] - self.textSize - 4
            cv2.rectangle(image, c1, c2, color, -1, cv2.LINE_AA)  # filled
            #cv2.putText(image, '{} {:.2f}'.format(cls_id, conf), (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
            image = self.cv2AddChineseText(image, '{} {:.2f}'.format(cls_id, conf), (c1[0], c1[1] - self.textSize - 2), (225, 255, 255), self.textSize)
        return image

    def detect(self, path, ext):
        path = path.replace('\\', '/')
        im = cv2.imread(path)
        im0, img = self.preprocess(im)

        pred = self.model(img, augment=False)[0]
        pred = pred.float()
        pred = non_max_suppression(pred, self.threshold, 0.5)

        pred_boxes = []
        image_info = []
        for det in pred:
            if det is not None and len(det):
                det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                for *x, conf, cls_id in det:
                    lbl = self.names[int(cls_id)]
                    x1, y1 = int(x[0]), int(x[1])
                    x2, y2 = int(x[2]), int(x[3])
                    pred_boxes.append((x1, y1, x2, y2, lbl, conf))
                    image_info.append([lbl, '{}×{}'.format(x2-x1, y2-y1), np.round(float(conf), 3)])

        im = self.plot_bboxes(im, pred_boxes)
        file_name = os.path.split(path)[1].split('.')[0]
        cv2.imwrite('./tmp/draw/{}.{}'.format(file_name, ext), im)

        return file_name + '.' + ext, image_info
